from numpy import array, sqrt, zeros, where, power, dot,  ones,  empty, inf
from numpy.linalg import inv, pinv, cond,  solve
from scipy.weave import converters
import scipy.weave as weave
from pyoptools.misc.Poly2D.Poly2D cimport *
#from pyoptools.misc.Poly2D import *

cimport numpy as np
import numpy as np
cimport cython

#from openopt import DFP


@cython.boundscheck(False)
@cython.wraparound(False)
def polyfit2d(x, y,z, int order=2):
    cdef int nc,nd,p
    nc= (order+2)*(order+1)/2 #_ord2i_(order)
    cdef np.ndarray[np.double_t, ndim=1, mode="c"] xa=array(x)
    cdef np.ndarray[np.double_t, ndim=1, mode="c"] ya=array(y)
    cdef np.ndarray[np.double_t, ndim=1, mode="c"] za=array(z)
    cdef np.ndarray[np.int_t, ndim=1, mode="c"] px,py
    
    px, py=i2pxpy(range(0, nc))
    nd=len(xa)
    cdef np.ndarray[np.double_t, ndim=2, mode="c"] XP=ones((nd, 2*(order+1)))
    cdef np.ndarray[np.double_t, ndim=2, mode="c"] YP=ones((nd, 2*(order+1)))
    for p in range(1, 2*(order+1)):
        XP[:, p]=XP[:, p-1]*xa
        YP[:, p]=YP[:, p-1]*ya

    #Calculating X and Y powers
    cdef np.ndarray[np.int_t, ndim=2, mode="c"] potx=empty((nc, nc),dtype=np.int)
    cdef np.ndarray[np.int_t, ndim=2, mode="c"] poty=empty((nc, nc),dtype=np.int)
    potx[:, :]=px
    potx=potx+potx.T
    
    poty[:, :]=py
    poty=poty+poty.T

    #Creating the Vandermonde matrix
    cdef np.ndarray[np.double_t, ndim=2, mode="c"] mat=empty((nc, nc))
    #The python code was changed to a C++ inline
    #for  ir in range(nc):
    #    for ic in range(nc):
    #        powx=potx[ir, ic]
    #        powy=poty[ir, ic]
    #        tv=XP[:, powx]*YP[:, powy]
    #        mat[ir, ic]=tv.sum()
    cdef int powx,powy,ir,ic, id
    cdef double sum
    for ir in range(nc):
        for ic in range(nc):
            powx=potx[ir,ic];
            powy=poty[ir,ic];
            sum=0;
            for id in range(nd): 
                sum=sum+XP[id,powx]*YP[id,powy]
            mat[ir,ic]=sum;
                    
    
    
    #imat=pinv(mat)

    cdef np.ndarray[np.double_t, ndim=2, mode="c"] vec=empty((nc, 1))
    #The python code was changed to C++ inline
    #for ic in range (nc):
    #    tv=XP[:, px[ic]]*YP[:, py[ic]]*z
    #    vec[ic, 0]=tv.sum()

    for ic in range(nc):
        sum=0;
        powx=px[ic]
        powy=py[ic]
        for id in range(nd):
            sum=sum+za[id]*XP[id,powx]*YP[id,powy]
        vec[ic,0]=sum;
    
    cohef= solve(mat, vec)
    ret_poly=poly2d(cohef[:,0])

    #Calculate error. Verify if this is the best way
    e= sqrt(power(array(z)-ret_poly.eval(xa, ya), 2).mean())
        
    return ret_poly, e
    
    
@cython.boundscheck(False)
@cython.wraparound(False)
def vander_matrix(x, y,z, int order=2):
    cdef int nc,nd,p
    nc= (order+2)*(order+1)/2 #_ord2i_(order)
    cdef np.ndarray[np.double_t, ndim=1, mode="c"] xa=array(x)
    cdef np.ndarray[np.double_t, ndim=1, mode="c"] ya=array(y)
    cdef np.ndarray[np.double_t, ndim=1, mode="c"] za=array(z)
    cdef np.ndarray[np.int_t, ndim=1, mode="c"] px,py
    
    px, py=i2pxpy(range(0, nc))
    nd=len(xa)
    cdef np.ndarray[np.double_t, ndim=2, mode="c"] XP=ones((nd, 2*(order+1)))
    cdef np.ndarray[np.double_t, ndim=2, mode="c"] YP=ones((nd, 2*(order+1)))
    for p in range(1, 2*(order+1)):
        XP[:, p]=XP[:, p-1]*xa
        YP[:, p]=YP[:, p-1]*ya

    #Calculating X and Y powers
    cdef np.ndarray[np.int_t, ndim=2, mode="c"] potx=empty((nc, nc),dtype=np.int)
    cdef np.ndarray[np.int_t, ndim=2, mode="c"] poty=empty((nc, nc),dtype=np.int)
    potx[:, :]=px
    potx=potx+potx.T
    
    poty[:, :]=py
    poty=poty+poty.T

    #Creating the Vandermonde matrix
    cdef np.ndarray[np.double_t, ndim=2, mode="c"] mat=empty((nc, nc))
    #The python code was changed to a C++ inline
    #for  ir in range(nc):
    #    for ic in range(nc):
    #        powx=potx[ir, ic]
    #        powy=poty[ir, ic]
    #        tv=XP[:, powx]*YP[:, powy]
    #        mat[ir, ic]=tv.sum()
    cdef int powx,powy,ir,ic, id
    cdef double sum
    for ir in range(nc):
        for ic in range(nc):
            powx=potx[ir,ic];
            powy=poty[ir,ic];
            sum=0;
            for id in range(nd): 
                sum=sum+XP[id,powx]*YP[id,powy]
            mat[ir,ic]=sum;
                    
    
    
    #imat=pinv(mat)

    cdef np.ndarray[np.double_t, ndim=2, mode="c"] vec=empty((nc, 1))
    #The python code was changed to C++ inline
    #for ic in range (nc):
    #    tv=XP[:, px[ic]]*YP[:, py[ic]]*z
    #    vec[ic, 0]=tv.sum()

    for ic in range(nc):
        sum=0;
        powx=px[ic]
        powy=py[ic]
        for id in range(nd):
            sum=sum+za[id]*XP[id,powx]*YP[id,powy]
        vec[ic,0]=sum;
    return mat,vec
    

#Local definitions for the optimized order 2 fitter

_potxo2=array([[ 0.,  1.,  0.,  2.,  1.,  0.], 
                [ 1.,  2.,  1.,  3.,  2.,  1.], 
                [ 0.,  1.,  0.,  2.,  1.,  0.], 
                [ 2.,  3.,  2.,  4.,  3.,  2.], 
                [ 1.,  2.,  1.,  3.,  2.,  1.], 
                [ 0.,  1.,  0.,  2.,  1.,  0.]])
_potyo2=array([[ 0.,  0.,  1.,  0.,  1.,  2.], 
                [ 0.,  0.,  1.,  0.,  1.,  2.], 
                [ 1.,  1.,  2.,  1.,  2.,  3.], 
                [ 0.,  0.,  1.,  0.,  1.,  2.], 
                [ 1.,  1.,  2.,  1.,  2.,  3.], 
                [ 2.,  2.,  3.,  2.,  3.,  4.]])
_pxo2, _pyo2=(array([0, 1, 0, 2, 1, 0]), array([0, 0, 1, 0, 1, 2]))

def polyfito2(x, y, z):
    """
    Polyfit function optimized for polynomials of order 2
    """
    nc= 6  #_ord2i_(order)
    xa=array(x)
    ya=array(y)
    px, py= _pxo2,_pyo2# _i2pxpy_(range(0, nc))
    nd=len(xa)
    XP=ones((nd, 6)) #ones((nd, 2*(order+1)))
    YP=ones((nd, 6))
    #for p in range(1, 6):
    #    XP[:, p]=XP[:, p-1]*xa
    #    YP[:, p]=YP[:, p-1]*ya
    cdef int p,id
    for p in range (1,6):
        for id in range(nd):
            XP[id,p]=XP[id,p-1]*xa[id];
            YP[id,p]=YP[id,p-1]*ya[id];

    #Calculating X and Y powers
    #potx=empty((nc, nc))
    #poty=empty((nc, nc))
    #potx[:, :]=px
    #potx=potx+potx.T
    #print "*",potx
    #poty[:, :]=py
    #poty=poty+poty.T
    #print "**", poty
    potx=_potxo2
    poty=_potyo2
    #Creating the Vandermonde matrix
    mat=empty((nc, nc))
    #The python code was changed to a C++ inline
    #for  ir in range(nc):
    #    for ic in range(nc):
    #        powx=potx[ir, ic]
    #        powy=poty[ir, ic]
    #        tv=XP[:, powx]*YP[:, powy]
    #        mat[ir, ic]=tv.sum()
    cdef int powx,powy
    cdef int ir,ic
    cdef double sum
    for ir in range[nc]:
        for ic in range (nc):
            powx=potx[ir,ic];
            powy=poty[ir,ic];
            sum=0;
            for id in range(nd):
                sum=sum+XP[id,powx]*YP[id,powy];
            mat[ir,ic]=sum;
        
    
    
    #imat=pinv(mat)

    vec=empty((nc, 1))
    #The python code was changed to C++ inline
    #for ic in range (nc):
    #    tv=XP[:, px[ic]]*YP[:, py[ic]]*z
    #    vec[ic, 0]=tv.sum()
    code="""
            int powx,powy;
            for(int ic=0;ic<nc;ic++)
            {
                double sum=0;
                powx=px(ic);
                powy=py(ic);
                for(int id=0;id<nd;id++)
                    sum=sum+z(id)*XP(id,powx)*YP(id,powy);
                vec(ic,0)=sum;
            }
    """
    err = weave.inline(code,['nc', 'nd', 'XP','YP','px', 'py', 'vec', 'z'],
                        type_converters=converters.blitz,
                        compiler = 'gcc')

    cohef= solve(mat, vec)
    ret_poly=poly2d(cohef, px=px, py=py)

    #Calculate error. Verify if this is the best way
    ep=cohef[0]+\
       cohef[1]*xa+\
       cohef[2]*ya+\
       cohef[3]*xa*xa+\
       cohef[4]*xa*ya+\
       cohef[5]*ya*ya
       
    e= sqrt(power(array(z)-ep, 2).mean())
        
    return ret_poly, e
_potxo1= array([[ 0.,   1.,   0.], 
          [ 1.,   2.,   1.], 
          [ 0.,   1.,   0.]])
_potyo1= array([[ 0.,   0.,   1.], 
          [ 0.,   0.,   1.], 
          [ 1.,   1.,   2.]])

_pxo1, _pyo1=(array([0, 1, 0]), array([0, 0, 1]))
def polyfito1(x, y, z):
    xa=array(x)
    ya=array(y)
    px, py=_pxo1, _pyo1
    nd=len(xa)
    XP=ones((nd, 4))
    YP=ones((nd, 4))
    code="""
            for(int p=1;p<4;p++)
            {
                for(int id=0;id<nd;id++)
                {
                    XP(id,p)=XP(id,p-1)*xa(id);
                    YP(id,p)=YP(id,p-1)*ya(id);
                }
            }
    """
    err = weave.inline(code,['nd', 'XP','YP','xa', 'ya'],
                        type_converters=converters.blitz,
                        compiler = 'gcc')

    #Calculating X and Y powers
    potx=_potxo1
    poty=_potyo1
    
    #Creating the Vandermonde matrix
    mat=empty((3, 3))

    code = """
            int powx,powy;
            for(int ir=0;ir<3;ir++)
                for(int ic=0;ic<3;ic++)
                    {   
                        powx=potx(ir,ic);
                        powy=poty(ir,ic);
                        double sum=0;
                        for(int id=0;id<nd;id++)
                            sum=sum+XP(id,powx)*YP(id,powy);
                        mat(ir,ic)=sum;
                    }
            """
    err = weave.inline(code,['potx', 'poty', 'XP', 'YP', 'mat', 'nd'],
                        type_converters=converters.blitz,
                        compiler = 'gcc')
    ######################
    
    
    #imat=pinv(mat)

    vec=empty((3, 1))

    code="""
            int powx,powy;
            for(int ic=0;ic<3;ic++)
            {
                double sum=0;
                powx=px(ic);
                powy=py(ic);
                for(int id=0;id<nd;id++)
                    sum=sum+z(id)*XP(id,powx)*YP(id,powy);
                vec(ic,0)=sum;
            }
    """
    err = weave.inline(code,[ 'nd', 'XP','YP','px', 'py', 'vec', 'z'],
                        type_converters=converters.blitz,
                        compiler = 'gcc')
    cohef= solve(mat, vec)
    ret_poly=poly2d(cohef)

    #Calculate error. Verify if this is the best way
    e= sqrt(power(array(z)-ret_poly.eval(xa, ya), 2).mean())
        
    return ret_poly, e




#~ def test(niter=10,Omax=4,  ndat=500):
    #~ from numpy.random import rand
    #~ for n in range(niter):
        #~ ncohef=ord2i(Omax)
        #~ cohef=100*rand(ncohef)
        #~ X=[] 
        #~ Y=[]
        #~ Z=[]
        #~ p=poly2d(cohef)
        #~ for j in range(ndat):
            #~ x=100.*rand()
            #~ y=100.*rand()
            #~ X.append(x)
            #~ Y.append(y)
            #~ Z.append(p.eval(x, y))
        #~ co, er=polyfit(X, Y, Z, Omax)
        #~ print cohef
        #~ print sqrt(power(cohef-co.cohef.T, 2).sum())/ncohef
        #~ print "*"
#~  Test using the open opt library it seems not to be needed
#~ def polyfitOO(x, y, z, order):
    #~ 
    #~ def f(C, D):
        #~ x=D[0]
        #~ y=D[1]
        #~ a=poly2d(cohef=C)
        #~ return a.eval(x, y)
    #~ ND=_ord2i_(order)
    #~ p, e=polyfit(array(x), array(y), array(z), order=order)
    #~ C=p.cohef
    #~ print C
    #~ D=array([x, y]).transpose()
    #~ Z=array(z).transpose()
    #~ lb = ND*[-inf]
    #~ ub = ND*[inf]
    #~ p = DFP(f, C , D, Z, lb=lb, ub=ub)
    #~ r = p.solve('nlp:ralg', plot=0, iprint = 10)
    #~ #print C
    #~ #print 'solution: '+str(r.xf)+'\n||residuals||^2 = '+str(r.ff)
    #~ ret_poly=poly2d(r.xf)
    #~ xa=array(x)
    #~ ya=array(y)
    #~ e= sqrt(power(array(z)-ret_poly.eval(xa, ya), 2).mean())
    #~ return ret_poly, e
